"""
@Author: Ling
@Date: 2022/5/12 12:31 
"""
import random
from matplotlib import pyplot as plt
import pandas as pd
import numpy as np
import seaborn as sns
from sklearn.linear_model import LinearRegression
from sklearn.linear_model import Ridge
from sklearn.preprocessing import PolynomialFeatures
from sklearn.metrics import mean_absolute_error


def train_test_split(data, state: str):
    return data[data['州'] == state].iloc[0:850], data[data['州'] == state].iloc[850:1034]


def train_test_split_add_random(data, state: str):
    train_data = data[data['州'] == state].iloc[0:850]
    train_data['确诊人数'] = train_data['确诊人数'] * (1 + 0.05 * np.random.rand(850) - 0.025)
    train_data['死亡人数'] = train_data['死亡人数'] * (1 + 0.05 * np.random.rand(850) - 0.025)
    test_data = data[data['州'] == state].iloc[850:1034]
    return train_data, test_data


class DataAnalysis():
    def __init__(self):
        self.data_preprocessing()
        pass

    def data_preprocessing(self):
        counties = pd.read_csv("dataset/NY-TIMES COVID-19 USA dataset (Regular updates)/States.csv")
        counties.columns = ['时间', '州', '联邦信息处理标准号', '确诊人数', '死亡人数']
        del counties['联邦信息处理标准号']
        # 构造新数据列
        # 利用“州”与“时间”列进行左外连接，以获取一个完全的各个州都以初始日期(2020-01-21)开始的新的序列
        # 在两个DataFrame中都构造一个同名新列，以进行后续处理
        city = pd.DataFrame(counties['州'].unique())
        time = pd.DataFrame(counties['时间'].unique())
        city['value'] = ''
        time['value'] = ''
        self.data = pd.merge(city, time, how='left', on='value')
        # 删除value列只留下有效数据
        del self.data['value']
        self.data.columns = ['州', '时间']
        # 初始化'确诊人数','死亡人数'两列并进行左外连接，后删除无效数据
        self.data[['确诊人数', '死亡人数']] = 0
        self.data = self.data.merge(counties, on=['州', '时间'], how='left')
        del self.data['确诊人数_x'], self.data['死亡人数_x']
        self.data.columns = ['州', '时间', '确诊人数', '死亡人数']
        # 用0填充值各州未发生疫情时的'确诊人数','死亡人数'两列数据
        self.data[['确诊人数', '死亡人数']] = self.data[['确诊人数', '死亡人数']].fillna(0)
        self.data[self.data['州'] == 'Washington']['时间'].count()  # 知晓每个时间序列一共1034个数据项
        self.data['州'].unique()  # 共有56个州
        # 数据预处理完成

    def paint_model_fit(self, column='死亡人数', state='Washington', model=1, degree=3, alpha=0.1, add_random=False):
        data = self.data[self.data['州'] == state]
        if add_random == True:
            train, test = train_test_split_add_random(self.data, state)
        else:
            train, test = train_test_split(self.data, state)
        if model == 1:
            y, MAE, model_name = ModelFit().model_1(train[column], test[column])
        elif model == 2:
            y, MAE, model_name = ModelFit().model_2(train[column], test[column], alpha)
        elif model == 3:
            y, MAE, model_name = ModelFit().model_3(train[column], test[column], degree)
        elif model == 4:
            y, MAE, model_name = ModelFit().model_4(train[column], test[column], degree, alpha)
        elif model == 5:
            y, MAE, model_name = ModelFit().model_5(train, test, column, alpha)
        elif model == 6:
            y, MAE, model_name = ModelFit().model_6(train[column], test[column])
        else:
            raise Exception("parameter error(model),No this parameter")
        # 取添加随机因素之后MAE最小的结果作为结果
        if add_random == True:
            train, test = train_test_split_add_random(self.data, state)
            for number in range(100):
                if model == 1:
                    the_y, the_MAE, model_name = ModelFit().model_1(train[column], test[column])
                elif model == 2:
                    the_y, the_MAE, model_name = ModelFit().model_2(train[column], test[column], alpha)
                elif model == 3:
                    the_y, the_MAE, model_name = ModelFit().model_3(train[column], test[column], degree)
                elif model == 4:
                    the_y, the_MAE, model_name = ModelFit().model_4(train[column], test[column], degree, alpha)
                elif model == 5:
                    the_y, the_MAE, model_name = ModelFit().model_5(train, test, column, alpha)
                elif model == 6:
                    the_y, the_MAE, model_name = ModelFit().model_6(train[column], test[column])
                else:
                    raise Exception("parameter error(model),No this parameter")
                if the_MAE < MAE:
                    y = the_y
                    MAE = the_MAE
        plt.figure()
        # 配色方案
        color = random.choice([[None, None, None, 'red'],
                               ['#70B3FF', '#776DFF', '#C700FF', '#8C4646'],
                               ['#F6AD00', '#F66900', '#E83200', '#98E800'],
                               ['#F4BD92', '#8D5894', '#38497F', '#77B3FF']])
        sns.set()
        sns.set_style('whitegrid')
        plt.subplots_adjust(top=1.5, right=1.5, wspace=0.5, hspace=0.5)
        plt.title(model_name + "  " + state + " 州 - " + column + "(截至2022年11月19日)", fontdict={'size': 20})
        plt.rcParams['font.sans-serif'] = ['SimHei']
        plt.rcParams['axes.unicode_minus'] = False
        plt.plot(range(850), y[0:850], label="模型拟合训练集", color=color[0], linewidth=2.0)
        plt.plot(range(850, 1034), y[850:1034], label="模型拟合测试集", color=color[1], linewidth=2.0)
        plt.plot(range(1034, 1200), y[1034:1200], label='模型预测结果', color=color[2], linewidth=2.0)
        plt.plot(range(1034), data[data['州'] == state][column], label=column + "实际值", linestyle='--', color=color[3],
                 linewidth=2.0)
        plt.xticks([0, 209, 419, 629, 839, 999, 1198],
                   ['2020-01-21', '2020-08-17', '2021-03-15', '2021-10-11', '2022-05-09', '2022-10-15', '2023-05-01'])
        plt.scatter(x=1033, y=y[1033], color='black', marker='s')
        plt.text(x=1033 - 50, y=y[1033] + 0.05 * y[1033], s='2022-11-19')
        yticks_obj = list(y[[0, 399, 839, 999]])
        yticks_obj.append(0)
        plt.yticks(yticks_obj)
        plt.text(800, y[0], "MAE:{:.2f}".format(float(MAE)), fontdict={'size': 18})
        plt.legend(loc=2)
        plt.xlabel('时间', fontdict={'size': 16})
        plt.ylabel(column, fontdict={'size': 16})
        plt.savefig('pic/{}_{}_{}_{}.png'.format(column, state, model, add_random), dpi=400, bbox_inches='tight')
        f = open('pic/{}_{}_{}_{}.png'.format(column, state, model, add_random), mode='rb')
        return f


class ModelFit():
    def model_1(self, train, test):  # 线性回归 train:850 test:90
        Linear_model = LinearRegression()
        Linear_model.fit(pd.DataFrame(range(850)), train)

        model_name = '一元线性回归'
        y = Linear_model.predict(X=pd.DataFrame(range(1200)))
        MAE = mean_absolute_error(test, y[850:1034])
        return y, MAE, model_name

    def model_2(self, train, test, alpha=0.1):  # 一元线性岭回归
        RidgeClassifier_model = Ridge(alpha=alpha)
        RidgeClassifier_model.fit(pd.DataFrame(range(850)), train)

        model_name = '一元线性岭回归'
        y = RidgeClassifier_model.predict(X=pd.DataFrame(range(1200)))
        MAE = mean_absolute_error(test, y[850:1034])
        return y, MAE, model_name

    def model_3(self, train, test, degree=3):  # 多次线性回归
        quadratic_featurizer = PolynomialFeatures(degree)  # 3次式时预测效果最好
        data_X2 = quadratic_featurizer.fit_transform(pd.DataFrame(range(1200)))
        Linear_model_x2 = LinearRegression()
        Linear_model_x2.fit(data_X2[0:850], train)

        model_name = '多次线性回归'
        y = Linear_model_x2.predict(data_X2)
        MAE = mean_absolute_error(test, y[850:1034])
        return y, MAE, model_name

    def model_4(self, train, test, degree=3, alpha=0.1):  # 多次项岭回归(最小二乘支持向量机)
        quadratic_featurizer = PolynomialFeatures(degree)  # 3次式时预测效果最好
        data_X2 = quadratic_featurizer.fit_transform(pd.DataFrame(range(1200)))
        RidgeClassifier_model_x2 = Ridge(alpha)
        RidgeClassifier_model_x2.fit(data_X2[0:850], train)

        model_name = '多次岭回归'
        y = RidgeClassifier_model_x2.predict(data_X2)
        MAE = mean_absolute_error(test, y[850:1034])
        return y, MAE, model_name

    def model_5(self, train, test, column, alpha=0.1):
        model_name = '相互干扰预测'
        if column == '死亡人数':
            ano_column = '确诊人数'
        else:
            ano_column = '死亡人数'

        # 取MAE最优多项式回归模型拟合结果
        MAE_y1 = None
        best_y1 = None
        for i in range(5):
            quadratic_featurizer_1 = PolynomialFeatures(i)
            data_X2_1 = quadratic_featurizer_1.fit_transform(pd.DataFrame(range(1200)))
            Linear_model_x2_1 = LinearRegression()
            Linear_model_x2_1.fit(data_X2_1[0:850], train[column])

            y1 = Linear_model_x2_1.predict(data_X2_1)
            if MAE_y1 is None:
                MAE_y1 = mean_absolute_error(test[column], y1[850:1034])
                best_y1 = y1
            elif MAE_y1 > mean_absolute_error(test[column], y1[850:1034]):
                MAE_y1 = mean_absolute_error(test[column], y1[850:1034])
                best_y1 = y1

        # 取MAE最优多项式回归模型拟合结果
        MAE_y2 = None
        best_y2 = None
        for i in range(5):
            quadratic_featurizer_2 = PolynomialFeatures(i)
            data_X2_2 = quadratic_featurizer_2.fit_transform(pd.DataFrame(range(1200)))
            Linear_model_x2_2 = LinearRegression()
            Linear_model_x2_2.fit(data_X2_2[0:850], train[ano_column])

            y2 = Linear_model_x2_2.predict(data_X2_2)
            if MAE_y2 is None:
                MAE_y2 = mean_absolute_error(test[column], y2[850:1034])
                best_y2 = y2
            elif MAE_y2 > mean_absolute_error(test[column], y2[850:1034]):
                MAE_y2 = mean_absolute_error(test[column], y2[850:1034])
                best_y2 = y2

        X = pd.DataFrame([np.block([0, best_y1[:-1]]), np.block([0, best_y2[:-1]])]).transpose().copy()

        RidgeClassifier_model_x2 = Ridge(alpha)
        RidgeClassifier_model_x2.fit(X[0:850], train[column])
        y = RidgeClassifier_model_x2.predict(X)
        MAE = mean_absolute_error(test[column], y[850:1034])

        return y, MAE, model_name

    def model_6(self, train, test):
        def GM11(x0):
            # 自定义灰色预测函数
            x1 = x0.cumsum()
            # 1-AGO序列
            z1 = (x1[:len(x1) - 1] + x1[1:]) / 2.0
            # 紧邻均值（MEAN）生成序列
            z1 = z1.reshape((len(z1), 1))
            B = np.append(-z1, np.ones_like(z1), axis=1)
            Yn = x0[1:].reshape((len(x0) - 1, 1))
            [[a], [b]] = np.dot(np.dot(np.linalg.inv(np.dot(B.T, B)), B.T), Yn)
            # 计算参数
            f = lambda k:(x0[0]-b/a)*np.exp(-a*(k-1))-(x0[0]-b/a)*np.exp(-a*(k-2))
            # 还原值
            return f

        Model_GM11 = GM11(train.to_numpy())
        model_name = '灰色预测'
        y = Model_GM11(np.arange(1200))
        MAE = mean_absolute_error(test, y[850:1034])
        return y, MAE, model_name
